import os
import subprocess
from loguru import logger
import importlib.util
import sys
import re # Import re for regex parsing of arguments
import json
from tools.tool_lifecycle_manager import ToolLifecycleManager

class ToolManager:
    def __init__(self):
        self.tools = {}
        self.lifecycle_manager = ToolLifecycleManager()
        self._load_tool_config()
        self._load_all_tools()

    def _load_tool_config(self):
        self.tool_config = {"band_all": False, "disabled_tools": []}
        config_path = os.path.join("tools", "tool_config.json")
        if os.path.exists(config_path):
            try:
                with open(config_path, 'r', encoding='utf-8') as f:
                    self.tool_config = json.load(f)
            except Exception as e:
                logger.error(f"Error loading tool_config.json: {e}")

    def _load_all_tools(self):
        self.tools = {} # Clear existing tools before reloading
        self._load_tool_descriptions() # Load descriptions first
        self._load_builtin_tools()
        self._load_generated_tools() # Load dynamically generated tools

    def _load_builtin_tools(self):
        builtin_tools_dir = os.path.join("tools", "builtin")
        if not os.path.exists(builtin_tools_dir):
            logger.warning(f"Builtin tools directory not found: {builtin_tools_dir}")
            return

        for filename in os.listdir(builtin_tools_dir):
            if filename.endswith(".py") and not filename.startswith("__"):
                # if "tools_func.py" in filename:
                #     continue # Skip tools_func.py
                module_name = f"tools.builtin.{filename[:-3]}"
                try:
                    # Use importlib.import_module for proper module loading
                    module = importlib.import_module(module_name)
                    for attr_name in dir(module):
                        attr = getattr(module, attr_name)
                        # Check if it's a callable function, not private, and defined in this module
                        if callable(attr) and not attr_name.startswith("_") and attr.__module__ == module.__name__:
                            tool_name = attr_name
                            if self.tool_config.get("band_all", False) and tool_name != "configure_tool_status":
                                continue
                            if tool_name in self.tool_config.get("disabled_tools", []):
                                continue
                            description = attr.__doc__.strip() if attr.__doc__ else f"Builtin tool: {tool_name}"
                            self.register_tool(tool_name, attr, description)
                            logger.info(f"Loaded builtin tool: {tool_name} from {filename}")
                except Exception as e:
                    logger.error(f"Error loading builtin tool from {filename}: {e}", exc_info=True)

        # Explicitly load tools from tools/builtin/tools_func.py
        # tools_func_path = os.path.join(builtin_tools_dir, "tools_func.py")
        # if os.path.exists(tools_func_path):
        #     spec = importlib.util.spec_from_file_location("tools_func", tools_func_path)
        #     if spec and spec.loader:
        #         tools_func_module = importlib.util.module_from_spec(spec)
        #         sys.modules["tools_func"] = tools_func_module
        #         try:
        #             spec.loader.exec_module(tools_func_module)
        #             for attr_name in dir(tools_func_module):
        #                 attr = getattr(tools_func_module, attr_name)
        #                 # Check if it's a callable function, not private, and defined in this module
        #                 if callable(attr) and not attr_name.startswith("_") and attr.__module__ == tools_func_module.__name__:
        #                     tool_name = attr_name
        #                     description = attr.__doc__.strip() if attr.__doc__ else f"Tool from tools_func.py: {tool_name}"
        #                     if tool_name in self.tools:
        #                         self.tools[tool_name]["func"] = attr
        #                         self.tools[tool_name]["description"] = description
        #                     else:
        #                         self.register_tool(tool_name, attr, description)
        #                     #logger.info(f"Loaded tool: {tool_name} from tools_func.py")
        #         except Exception as e:
        #             logger.error(f"Error loading tools from tools_func.py: {e}", exc_info=True)
        #     else:
        #         logger.error(f"Could not get spec or loader for tools_func.py")
        # else:
        #     logger.warning(f"tools_func.py not found at {tools_func_path}")

    def _load_tool_descriptions(self):
        tool_info_path = os.path.join("tools", "tools_info.md")
        if not os.path.exists(tool_info_path):
            logger.warning(f"Tool info file not found: {tool_info_path}. Using hardcoded descriptions.")
            return

        try:
            with open(tool_info_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            tool_sections = content.split("### ")[1:]
            for section in tool_sections:
                lines = section.strip().split('\n')
                tool_name_raw = lines[0].strip()
                tool_name = tool_name_raw.lower() # Store lowercase for internal lookup

                description = ""
                usage = ""
                arguments = [] # List to store parsed arguments
                
                arg_section_started = False

                for line in lines[1:]:
                    line = line.strip()
                    if line.startswith("**Description:**"):
                        description = line.replace("**Description:**", "").strip()
                    elif line.startswith("**Usage:**"):
                        usage = line.replace("**Usage:**", "").strip()
                    elif line.startswith("**Arguments:**"):
                        arg_section_started = True
                    elif arg_section_started and line.startswith("- {"):
                        # Parse argument line: - {type} param_name: Description
                        match = re.match(r'^- {([^}]+)}\s*([^:]+):\s*(.*)$', line)
                        if match:
                            arg_type = match.group(1).strip()
                            arg_name = match.group(2).strip()
                            arg_desc = match.group(3).strip()
                            
                            # Filter out userid and openid
                            if arg_name.lower() not in ["userid", "openid"]:
                                arguments.append({"name": arg_name, "type": arg_type, "description": arg_desc})
                    elif arg_section_started and not line: # End of argument section
                        arg_section_started = False
                    elif not arg_section_started and line: # Other lines before arguments
                        # This part was previously concatenating all lines into full_description
                        # Now we are more structured, so we only care about specific tags
                        pass

                if tool_name in self.tools:
                    self.tools[tool_name]["description"] = description
                    self.tools[tool_name]["usage"] = usage
                    self.tools[tool_name]["arguments"] = arguments
                    self.tools[tool_name]["original_name"] = tool_name_raw # Store original name for display
                else:
                    # If tool is not registered by _load_builtin_tools or _load_generated_tools, register it now
                    # This handles cases where tools_info.md might describe tools not yet loaded from code
                    self.tools[tool_name] = {
                        "func": None, # No function object yet, will be filled if found later
                        "description": description,
                        "usage": usage,
                        "arguments": arguments,
                        "original_name": tool_name_raw
                    }
                    logger.debug(f"Tool '{tool_name}' found in tools_info.md but not registered by code. Registering with description only.")

            #logger.info(f"Successfully loaded tool descriptions from {tool_info_path}")
        except Exception as e:
            logger.error(f"Error loading tool descriptions from {tool_info_path}: {e}", exc_info=True)

    def get_tool_description(self, tool_name: str) -> str:
        # This method might need to be updated to return a more structured description
        # For now, it returns the simple description
        return self.tools.get(tool_name, {}).get("description", "No description available.")

    def get_tool_func(self, tool_name: str):
        return self.tools.get(tool_name, {}).get("func")

    def get_all_tool_descriptions(self, allowed_tools: list[str] | None = None) -> str:
        descriptions = []
        for name, tool_info in self.tools.items():
            original_name = tool_info.get("original_name", name) # Use original name for display
            
            if self.tool_config.get("band_all", False) and original_name != "configure_tool_status":
                continue
            if original_name in self.tool_config.get("disabled_tools", []):
                continue

            # Apply filtering if allowed_tools list is provided
            if allowed_tools is not None and original_name not in allowed_tools:
                continue

            description = tool_info.get("description", "No description available.")
            usage = tool_info.get("usage", "")
            arguments = tool_info.get("arguments", [])
            if original_name == "quote" or original_name == "urlencode" or original_name == "LLMInterface":
                continue

            tool_desc_str = f"- {original_name}: {description}"
            if usage:
                tool_desc_str += f"\n  Usage: {usage}"
            if arguments:
                tool_desc_str += "\n  Arguments:"
                for arg in arguments:
                    tool_desc_str += f"\n    - {arg['name']} ({arg['type']}): {arg['description']}"
                
                # Generate an example for params XML
                params_example = ""
                if len(arguments) == 1 and arguments[0]['name'] == 'query' and original_name == 'web_search':
                    params_example = f"<params><query>example_query</query></params>"
                elif len(arguments) == 1:
                    # If only one argument, use <value> tag as per rootAgent.py example
                    params_example = f"<params><value>example_{arguments[0]['name']}</value></params>"
                elif len(arguments) > 1:
                    params_tags = "".join([f"<{arg['name']}>example_{arg['name']}</{arg['name']}>" for arg in arguments])
                    params_example = f"<params>{params_tags}</params>"
                
                if params_example:
                    tool_desc_str += f"\n  Example Params XML: {params_example}"

            descriptions.append(tool_desc_str)
        return "\n".join(descriptions)

    def _load_generated_tools(self):
        generated_tools_dir = os.path.join("tools", "generated")
        os.makedirs(generated_tools_dir, exist_ok=True)
        
        # Add the generated tools directory to sys.path temporarily
        if generated_tools_dir not in sys.path:
            sys.path.insert(0, generated_tools_dir)

        for filename in os.listdir(generated_tools_dir):
            if filename.endswith(".py") and filename != "__init__.py":
                module_name = filename[:-3] # Remove .py extension
                try:
                    # Import the module dynamically
                    module = __import__(module_name)
                    # Iterate through module's attributes to find functions
                    for attr_name in dir(module):
                        attr = getattr(module, attr_name)
                        if callable(attr) and not attr_name.startswith("_") and attr_name != module_name and attr.__module__ == module.__name__: # Avoid private and module-level functions
                            # Assuming the function name is the tool name
                            tool_name = attr_name
                            if self.tool_config.get("band_all", False) and tool_name != "configure_tool_status":
                                continue
                            if tool_name in self.tool_config.get("disabled_tools", []):
                                continue
                            # Attempt to get docstring as description
                            description = attr.__doc__.strip() if attr.__doc__ else f"Dynamically generated tool: {tool_name}"
                            if tool_name in self.tools:
                                self.tools[tool_name]["func"] = attr
                                self.tools[tool_name]["description"] = description
                            else:
                                self.register_tool(tool_name, attr, description)
                            logger.info(f"Dynamically loaded tool: {tool_name} from {filename}")
                except Exception as e:
                    logger.error(f"Error loading generated tool from {filename}: {e}", exc_info=True)
        
        # Remove the generated tools directory from sys.path
        if generated_tools_dir in sys.path:
            sys.path.remove(generated_tools_dir)

    def register_tool(self, tool_name: str, func, description: str = ""):
        if tool_name in self.tools:
            # If tool already exists, update its function and description
            self.tools[tool_name]["func"] = func
            self.tools[tool_name]["description"] = description
            logger.info(f"Tool '{tool_name}' updated successfully.")
        else:
            self.tools[tool_name] = {"func": func, "description": description}
            logger.info(f"Tool '{tool_name}' registered successfully.")

    def register_new_tool(self, tool_name: str, func, description: str, usage: str, arguments: list):
        """
        Registers a newly created tool and updates tools_info.md.
        """
        self.register_tool(tool_name, func, description)
        
        tool_info_path = os.path.join("tools", "tools_info.md")
        with open(tool_info_path, 'a', encoding='utf-8') as f:
            f.write(f"\n### {tool_name}\n")
            f.write(f"**Description:** {description}\n")
            f.write(f"**Usage:** {usage}\n")
            if arguments:
                f.write("**Arguments:**\n")
                for arg in arguments:
                    f.write(f"- {{{arg['type']}}} {arg['name']}: {arg['description']}\n")
            f.write("\n")
        logger.info(f"New tool '{tool_name}' registered and added to tools_info.md.")
        self._load_tool_descriptions() # Reload descriptions to include the new tool

    def promote_tool(self, tool_name: str) -> bool:
        success = self.lifecycle_manager.promote_tool(tool_name)
        if success:
            self._load_all_tools() # Reload tools after promotion
        return success

    def deprecate_tool(self, tool_name: str) -> bool:
        success = self.lifecycle_manager.deprecate_tool(tool_name)
        if success:
            self._load_all_tools() # Reload tools after deprecation
        return success